package com.zcmyth.tree;

import java.util.ArrayDeque;
import java.util.Queue;
import java.util.Stack;

public class BinaryTree<T> {
  private final Node<T> root;

  public BinaryTree(Node<T> root) {
    this.root = root;
  }

  public Node<T> root() {
    return root;
  }

  /**
   * Visits {@link Node} so we can change the {@code visit} mark.
   */
  public void bfs(TreeVisitor<Node<T>> visitor) {
    Queue<Node<T>> queue = new ArrayDeque<Node<T>>();
    queue.add(root);
    while (!queue.isEmpty()) {
      Node<T> element = queue.poll();
      visitor.visit(element);
      if (element.getLeft() != null) {
        queue.add(element.getLeft());
      }
      if (element.getRight() != null) {
        queue.add(element.getRight());
      }
    }
  }

  public void clear() {
    bfs(new TreeVisitor<Node<T>>() {
      @Override
      public void visit(Node<T> e) {
        e.setVisited(false);
      }
    });
  }

  /**
   * Pre-order doesn't need visit mark.
   * 
   * @param visitor
   */
  public void preOrder(TreeVisitor<T> visitor) {
    Stack<Node<T>> stack = new Stack<Node<T>>();
    stack.push(root);
    while (!stack.isEmpty()) {
      Node<T> element = stack.pop();
      visitor.visit(element.getData());
      if (element.getRight() != null) {
        stack.push(element.getRight());
      }
      if (element.getLeft() != null) {
        stack.push(element.getLeft());
      }
    }
  }

  /**
   * In-order need visit mark, clear first.
   */
  public void inOrder(TreeVisitor<T> visitor) {
    clear();
    Stack<Node<T>> stack = new Stack<Node<T>>();
    stack.push(root);
    while (!stack.isEmpty()) {
      Node<T> element = stack.peek();
      if(element.getLeft() == null || element.getLeft().isVisited()) {
        stack.pop();
        visitor.visit(element.getData());
        element.setVisited(true);
        if(element.getRight() != null) {
          stack.push(element.getRight());
        }
      } else {
        stack.push(element.getLeft());
      }
    }
  }

  public void postOrder(TreeVisitor<T> visitor) {
    clear();
    Stack<Node<T>> stack = new Stack<Node<T>>();
    stack.push(root);
    while (!stack.isEmpty()) {
      Node<T> element = stack.peek();
      if(element.getLeft() != null && (!element.getLeft().isVisited())) {
        stack.push(element.getLeft());
      } else if(element.getRight() != null && (!element.getRight().isVisited())){
        stack.push(element.getRight());
      } else {
        visitor.visit(element.getData());
        stack.pop();
        element.setVisited(true);
      }
    }
  }

  public static class Node<E> {

    private final E data;
    private Node<E> left;
    private Node<E> right;

    private boolean visited;

    public Node(E d) {
      this.data = d;
      this.visited = false;
    }

    public E getData() {
      return data;
    }

    public Node<E> getLeft() {
      return left;
    }

    public Node<E> getRight() {
      return right;
    }

    public void setLeft(Node<E> n) {
      left = n;
    }

    public void setRight(Node<E> n) {
      right = n;
    }

    public void setVisited(boolean visited) {
      this.visited = visited;
    }

    public boolean isVisited() {
      return visited;
    }
  }
}
